#include <sys/epoll.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <unistd.h>
#include <fcntl.h>
#include <iostream>
#include <vector>
#include <string> // Include string header
#include <cstring>
#include <chrono>
#include <map>    // Include map header
#include <cerrno> // Include errno header

#define MAX_EVENTS 1024
#define BUFFER_SIZE 2048

// Structure to hold connection state (including write buffer)
struct ConnectionInfo {
    std::string write_buffer;
};

// Global map to store connection info (including write buffers) per fd
std::map<int, ConnectionInfo> client_connections;

// Function to set socket non-blocking
int set_nonblocking(int fd) {
    int flags = fcntl(fd, F_GETFL, 0);
    if (flags == -1) {
        perror("fcntl(F_GETFL)");
        return -1;
    }
    return fcntl(fd, F_SETFL, flags | O_NONBLOCK);
}

// Function to handle closing a client connection
void close_client(int fd, int epoll_fd) {
    epoll_ctl(epoll_fd, EPOLL_CTL_DEL, fd, nullptr); // Remove from epoll
    close(fd);                                       // Close the socket
    client_connections.erase(fd);                    // Remove connection state
    std::cout << "Client disconnected: fd=" << fd << std::endl;
}

// --- Handle Read ---
void handle_read(int fd, int epoll_fd) {
    // auto tp0 = std::chrono::high_resolution_clock::now();
    char buf[BUFFER_SIZE];
    std::string received_data;
    bool connection_closed = false;
    bool error_occurred = false;

    // Read data in a loop due to Edge Triggering (EPOLLET)
    while (true) {
        ssize_t bytes_read = read(fd, buf, sizeof(buf));

        if (bytes_read > 0) {
            received_data.append(buf, bytes_read);
        }
        else if (bytes_read == 0) {
            // Connection closed by peer
            connection_closed = true;
            break;
        }
        else if (bytes_read == -1 && (errno == EAGAIN || errno == EWOULDBLOCK)) {
            // No more data to read for now (EAGAIN/EWOULDBLOCK for non-blocking sockets)
            break;
        }
        else {
            // Actual read error
            perror("read");
            error_occurred = true;
            break;
        }

    }

    if (connection_closed || error_occurred) {
        close_client(fd, epoll_fd);
        return;
    }

    // If we received data, prepare the response and register for write
    if (!received_data.empty()) {
        std::cout << "Received " << received_data.length() << " bytes from fd=" << fd << std::endl;

        // Prepare the HTTP response
        std::string response = "HTTP/1.1 200 OK\r\nContent-Length: " +
            std::to_string(received_data.size()) +
            "\r\nConnection: keep-alive" + // Keep connection open
            "\r\n\r\n" + received_data;

        // Store the response in the connection's write buffer
        // If data already exists (unlikely for simple echo but good practice), append
        client_connections[fd].write_buffer += response;

        // Modify epoll registration to include EPOLLOUT
        epoll_event event{};
        event.events = EPOLLIN | EPOLLOUT | EPOLLET; // Listen for read and write
        event.data.fd = fd;
        if (epoll_ctl(epoll_fd, EPOLL_CTL_MOD, fd, &event) == -1) {
            perror("epoll_ctl (MOD to add EPOLLOUT)");
            close_client(fd, epoll_fd); // Close on error
        }
        // std::cout << "Registered fd=" << fd << " for EPOLLOUT" << std::endl;
    }
    // else: No data read this time, might happen with ET if previous read emptied buffer
    //       but client didn't close. Nothing to do.

    // auto tp1 = std::chrono::high_resolution_clock::now();
    // auto duration1 = std::chrono::duration_cast<std::chrono::microseconds>(tp1 - tp0).count();
    // std::cout << "Read handling time for fd=" << fd << ": " << duration1 << "us" << std::endl;
}


// --- Handle Write ---
void handle_write(int fd, int epoll_fd) {
    // auto tp0 = std::chrono::high_resolution_clock::now();
    std::string& data_to_write = client_connections[fd].write_buffer; // Get reference to buffer
    bool error_occurred = false;

    // Write data in a loop until buffer is empty or write blocks
    while (!data_to_write.empty()) {
        ssize_t bytes_written = write(fd, data_to_write.c_str(), data_to_write.size());

        if (bytes_written >= 0) {
            // std::cout << "Wrote " << bytes_written << " bytes to fd=" << fd << std::endl;
            data_to_write.erase(0, bytes_written); // Remove written part from buffer
            if (data_to_write.empty()) {
                // std::cout << "Write buffer empty for fd=" << fd << std::endl;
                break; // Exit loop if buffer is now empty
            }
            // Continue loop if more data remains
        }
        else { // bytes_written == -1
            if (errno == EAGAIN || errno == EWOULDBLOCK) {
                // Socket buffer is full, cannot write more now.
                // EPOLLOUT will trigger again when space is available.
                // std::cout << "Write blocked (EAGAIN) for fd=" << fd << std::endl;
                break; // Exit the write loop for now
            }
            else {
                // Actual write error
                perror("write");
                error_occurred = true;
                break;
            }
        }
    }

    if (error_occurred) {
        close_client(fd, epoll_fd);
        return;
    }

    // If the buffer is now empty, stop monitoring EPOLLOUT
    if (data_to_write.empty()) {
        epoll_event event{};
        event.events = EPOLLIN | EPOLLET; // Only listen for read events now
        event.data.fd = fd;
        if (epoll_ctl(epoll_fd, EPOLL_CTL_MOD, fd, &event) == -1) {
            perror("epoll_ctl (MOD to remove EPOLLOUT)");
            close_client(fd, epoll_fd); // Close on error
        }
        // std::cout << "De-registered fd=" << fd << " from EPOLLOUT" << std::endl;
    }

    // auto tp1 = std::chrono::high_resolution_clock::now();
    // auto duration1 = std::chrono::duration_cast<std::chrono::microseconds>(tp1 - tp0).count();
    // std::cout << "Write handling time for fd=" << fd << ": " << duration1 << "us" << std::endl;
}

int main() {
    // --- Socket Creation, Bind, Listen (same as before) ---
    int listen_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (listen_fd == -1) {
        perror("socket");
        return 1;
    }

    int opt = 1;
    if (setsockopt(listen_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) == -1) {
        perror("setsockopt(SO_REUSEADDR)");
        close(listen_fd);
        return 1;
    }

    // Set listener non-blocking *before* epoll add is generally good practice
    if (set_nonblocking(listen_fd) == -1) {
        close(listen_fd);
        return 1;
    }


    sockaddr_in addr{};
    addr.sin_family = AF_INET;
    addr.sin_port = htons(8080);
    addr.sin_addr.s_addr = INADDR_ANY;

    if (bind(listen_fd, (sockaddr*)&addr, sizeof(addr)) == -1) {
        perror("bind");
        close(listen_fd);
        return 1;
    }

    if (listen(listen_fd, SOMAXCONN) == -1) {
        perror("listen");
        close(listen_fd);
        return 1;
    }

    // --- Epoll Creation and Setup ---
    int epoll_fd = epoll_create1(0);
    if (epoll_fd == -1) {
        perror("epoll_create1");
        close(listen_fd);
        return 1;
    }

    epoll_event event{};
    event.events = EPOLLIN; // Listen for incoming connections
    event.data.fd = listen_fd;
    if (epoll_ctl(epoll_fd, EPOLL_CTL_ADD, listen_fd, &event) == -1) {
        perror("epoll_ctl (listen_fd)");
        close(listen_fd);
        close(epoll_fd);
        return 1;
    }

    std::vector<epoll_event> events(MAX_EVENTS);

    std::cout << "Echo server started on port 1127..." << std::endl;

    // --- Event Loop ---
    while (true) {
        int n = epoll_wait(epoll_fd, events.data(), MAX_EVENTS, -1);
        if (n == -1) {
            // EINTR is an interrupted system call, can often be ignored/retried
            if (errno == EINTR) {
                continue;
            }
            perror("epoll_wait");
            break; // Exit loop on other errors
        }

        for (int i = 0; i < n; ++i) {
            int current_fd = events[i].data.fd;
            uint32_t current_events = events[i].events;

            // --- Handle New Connections ---
            if (current_fd == listen_fd) {
                // Loop to accept all pending connections (due to potential level-triggering nature of accept)
                // Though epoll often behaves edge-triggered for accept too, this is safer.
                while (true) {
                    sockaddr_in client_addr{};
                    socklen_t addr_len = sizeof(client_addr);
                    int client_fd = accept(listen_fd, (sockaddr*)&client_addr, &addr_len);

                    if (client_fd == -1) {
                        if (errno == EAGAIN || errno == EWOULDBLOCK) {
                            // No more incoming connections for now
                            break;
                        }
                        else {
                            perror("accept");
                            break; // Error accepting
                        }
                    }

                    std::cout << "New connection accepted: fd=" << client_fd << std::endl;

                    if (set_nonblocking(client_fd) == -1) {
                        close(client_fd);
                        continue; // Handle next event/connection
                    }

                    // Add new client to epoll for read events initially
                    epoll_event client_event{};
                    client_event.events = EPOLLIN | EPOLLET; // Edge-triggered read
                    client_event.data.fd = client_fd;
                    if (epoll_ctl(epoll_fd, EPOLL_CTL_ADD, client_fd, &client_event) == -1) {
                        perror("epoll_ctl (client_fd)");
                        close(client_fd);
                    }
                    else {
                        // Initialize connection state
                        client_connections[client_fd] = ConnectionInfo{};
                    }
                } // End accept loop
            }
            // --- Handle Existing Connections ---
            else {
                // Check for errors first (EPOLLERR, EPOLLHUP)
                // Note: EPOLLIN/EPOLLOUT might be set alongside errors
                if (current_events & (EPOLLERR | EPOLLHUP)) {
                    std::cerr << "Epoll error or hangup on fd=" << current_fd << ". Events: " << current_events << std::endl;
                    close_client(current_fd, epoll_fd);
                    continue; // Process next event
                }

                // Handle Write Event (if socket is writable and we have data)
                if (current_events & EPOLLOUT) {
                    // Check if we actually have data to write before calling handler
                    // (Defensive check, handle_write also checks)
                    if (client_connections.count(current_fd) && !client_connections[current_fd].write_buffer.empty()) {
                        handle_write(current_fd, epoll_fd);
                    }
                    else {
                        // Spurious EPOLLOUT? Or logic error? For now, just log.
                        // std::cout << "EPOLLOUT received for fd=" << current_fd << " but no data in buffer." << std::endl;
                        // Optionally, remove EPOLLOUT monitoring here if it wasn't removed properly before.
                        epoll_event mod_event{};
                        mod_event.events = EPOLLIN | EPOLLET;
                        mod_event.data.fd = current_fd;
                        epoll_ctl(epoll_fd, EPOLL_CTL_MOD, current_fd, &mod_event);
                    }
                }

                // Handle Read Event (must check *after* potential close in write handler)
                // Check if the connection still exists before reading
                if (client_connections.count(current_fd) && (current_events & EPOLLIN)) {
                    handle_read(current_fd, epoll_fd);
                }

            } // End handling existing connection
        } // End loop through events
    } // End main loop

    // --- Cleanup ---
    std::cout << "Shutting down server..." << std::endl;
    // Close all remaining client connections
    for (auto const& [fd, val] : client_connections) {
        close(fd);
    }
    client_connections.clear();
    close(listen_fd);
    close(epoll_fd);
    return 0;
}